import cv2
import json
import os
import numpy as np
import matplotlib as plt
from sklearn.cluster import KMeans
json_dir = '/data/lzy/tile_round1_train_20201231/train_annos.json'
high = []
wid = []
size = []
hw_ratio = []
classes = []
with open(json_dir) as f:
    annos=json.load(f)

for i in range(len(annos)):
    anno = annos[i]
    name = anno['name']
    bbox = anno['bbox']
    x1 = [bbox[0], bbox[1]]
    x2 = [bbox[2], bbox[3]]
    height = abs(x1[1] - x2[1])
    width = abs(x1[0] - x2[0])
    cate = anno['category']
    classes.append(cate)
    size_now = height * width
    hw_r = height/width
    high.append(height)
    wid.append(width)
    size.append(size_now)
    hw_ratio.append(hw_r)
mdl_size = KMeans(n_clusters=10, random_state=2021)
mdl_ratio = KMeans(n_clusters=20, random_state=2021)
np_high = np.asarray(high)
np_wid = np.asarray(wid)
np_size = np.asarray(size).reshape((-1, 1))
np_hw = np.asarray(hw_ratio).reshape((-1, 1))
np_classes = np.asarray(classes)
clf_size = mdl_size.fit(np_size)
clf_ratio = mdl_ratio.fit(np_hw)
np.save('high.npy', np_high)
np.save('width.npy', np_wid)
np.save('size.npy', np_size)
np.save('hw_ratio.npy', np_hw)
np.save('np_classes.npy', np_classes)